#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import random
import shutil

sys.path.append('.')

"""
   生成人脸ReID数据集
"""


def read_file_list(file_list, mode=0):
    faceid_subfaceid_files = {}
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read File Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            if mode == 0:
                # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
                lines = line.strip().split()
                faceid = lines[0]
                subfaceid = lines[1]
                file_name = lines[2]
                faceid_subfaceid = '{} {}'.format(faceid, subfaceid)
            elif mode == 1:
                # gd-zs-sqqjd-dxssjyeq-150100414a54443452064d2742c73600-video-20190913_225440/0001/0001_F0001B00.jpg
                line = line.strip()
                lines = line.split('/')
                file_name = line
                faceid_subfaceid = '{} {}'.format(lines[0], lines[1])

            file_selects = faceid_subfaceid_files.get(faceid_subfaceid, None)
            if file_selects is None:
                faceid_subfaceid_files[faceid_subfaceid] = [file_name]
            else:
                file_selects.append(file_name)

        print('Read File Count: {}'.format(file_count))
    return faceid_subfaceid_files


def copy_file(image_dir, files, pid, out_list_file, out_image_path, out_pair_list_file):
    dst_dir = os.path.join(out_image_path, pid)
    if not os.path.exists(dst_dir):
        os.makedirs(dst_dir)
        os.system('chmod a+wr {}'.format(dst_dir))
    random.shuffle(files)
    for fr_num, file_name in enumerate(files):
        src_file_path = os.path.join(image_dir, file_name)
        dst_file_name = '{}_F{:04d}B00.jpg'.format(pid, fr_num)
        dst_file_path = os.path.join(dst_dir, dst_file_name)
        shutil.copy(src_file_path, dst_file_path)
        os.system('chmod a+wr {}'.format(dst_file_path))
        out_list_file.write(os.path.join(pid, dst_file_name) + '\n')
        pair_info = '{}:{}\n'.format(file_name, dst_file_name)
        out_pair_list_file.write(pair_info)


def main():
    parser = argparse.ArgumentParser(description="Replace Face ReID Image")
    parser.add_argument(
        "--image_dir", default="/rootfs/media/kasim/Data/data/FaceImage/Image", help="path to image file list", type=str
    )
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data/data/FaceImage/filter100/FaceRecSel.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data/data/ErisedFaceReID", help="path to output", type=str
    )
    parser.add_argument('--test_ratio', type=float, default=0.2, help='test ratio')
    parser.add_argument('--seed', type=int, default=123456, help='random seed')
    parser.add_argument('--mode', type=int, default=0, help='file list mode')
    args = parser.parse_args()

    random.seed(args.seed)

    image_dir = args.image_dir
    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    faceid_files = read_file_list(args.file_list, args.mode)

    id_list = list(faceid_files.keys())
    random.shuffle(id_list)
    test_count = int(round(len(id_list) * args.test_ratio))
    if test_count > len(id_list):
        test_count = len(id_list)

    test_id_list = id_list[0:test_count]
    if test_count >= len(id_list):
        train_id_list = []
    else:
        train_id_list = id_list[test_count:]

    train_id_count = 0
    train_image_count = 0
    if len(train_id_list) > 0:
        train_image_path = os.path.join(output_dir, 'train')
        train_list_path = os.path.join(output_dir, 'train.txt')
        train_pair_files_path = os.path.join(output_dir, 'train_pair_files.txt')
        train_pair_ids_path = os.path.join(output_dir, 'train_pair_ids.txt')
        train_pair_files_file = open(train_pair_files_path, 'w')
        train_pair_ids_file = open(train_pair_ids_path, 'w')
        with open(train_list_path, 'w') as file:
            for i, faceid in enumerate(train_id_list):
                files = faceid_files[faceid]
                if len(files) > 0:
                    train_id_count += 1
                    train_image_count += len(files)
                    pid = '{:04d}'.format(i + 1)
                    pair_info = '{}:{}\n'.format(faceid, pid)
                    train_pair_ids_file.write(pair_info)
                    copy_file(image_dir, files, pid, file, train_image_path, train_pair_files_file)
                if (i + 1) % 100 == 0:
                    print('Copy Train Face ID Count: {}/{}'.format(i+1, len(train_id_list)))
        train_pair_files_file.close()
        train_pair_ids_file.close()
        os.system('chmod a+wr {}'.format(train_image_path))

    query_id_count = 0
    query_image_count = 0
    gallery_id_count = 0
    gallery_image_count = 0
    if len(test_id_list) > 0:
        query_image_path = os.path.join(output_dir, 'query')
        query_list_path = os.path.join(output_dir, 'query.txt')
        gallery_image_path = os.path.join(output_dir, 'gallery')
        gallery_list_path = os.path.join(output_dir, 'gallery.txt')
        query_file = open(query_list_path, 'w')
        gallery_file = open(gallery_list_path, 'w')

        query_pair_files_path = os.path.join(output_dir, 'query_pair_files.txt')
        query_pair_ids_path = os.path.join(output_dir, 'query_pair_ids.txt')
        query_pair_files_file = open(query_pair_files_path, 'w')
        query_pair_ids_file = open(query_pair_ids_path, 'w')

        gallery_pair_files_path = os.path.join(output_dir, 'gallery_pair_files.txt')
        gallery_pair_ids_path = os.path.join(output_dir, 'gallery_pair_ids.txt')
        gallery_pair_files_file = open(gallery_pair_files_path, 'w')
        gallery_pair_ids_file = open(gallery_pair_ids_path, 'w')

        for i, faceid in enumerate(test_id_list):
            files = faceid_files[faceid]
            if len(files) > 0:
                pid = '{:04d}'.format(i + 1)
                random.shuffle(files)
                query_count = len(files) // 2
                query_files = files[0:query_count]
                gallery_files = files[query_count:]
                if len(query_files) > 0:
                    query_id_count += 1
                    query_image_count += len(query_files)
                    pair_info = '{}:{}\n'.format(faceid, pid)
                    query_pair_ids_file.write(pair_info)
                    copy_file(image_dir, query_files, pid, query_file, query_image_path, query_pair_files_file)
                if len(gallery_files) > 0:
                    gallery_id_count += 1
                    gallery_image_count += len(gallery_files)
                    pair_info = '{}:{}\n'.format(faceid, pid)
                    gallery_pair_ids_file.write(pair_info)
                    copy_file(image_dir, gallery_files, pid, gallery_file, gallery_image_path, gallery_pair_files_file)
            if (i + 1) % 100 == 0:
                print('Copy Test Face ID Count: {}/{}'.format(i+1, len(test_id_list)))
        gallery_pair_files_file.close()
        gallery_pair_ids_file.close()
        query_pair_files_file.close()
        query_pair_ids_file.close()

        query_file.close()
        gallery_file.close()
        os.system('chmod a+wr {}'.format(gallery_image_path))
        os.system('chmod a+wr {}'.format(query_image_path))
    os.system('chmod a+wr {}/*.txt'.format(output_dir))

    print('Train ID Count: {}, Train Image Count: {}'.format(train_id_count, train_image_count))
    print('Query ID Count: {}, Query Image Count: {}'.format(query_id_count, query_image_count))
    print('Gallery ID Count: {}, Gallery Image Count: {}'.format(gallery_id_count, gallery_image_count))

    print('finish!')


if __name__ == '__main__':
    main()
